#pragma once

#include "libshmpy.hpp"

namespace shmpy::mem::core {
    /*
     * Trunk is where to hold the same size of memory block. e.g.
     * a 32 bytes Trunk will hold 1024 bytes of block. therefore, its
     * total size is 32kb; Trunk Meta is implemented to track a trunk's status.
     * Since this is design for shared memory, instead of using pointer to indicate
     * next or last trunk location, pshift_t is used to calculate their address when
     * a Shared Memory Resource is attach to a process.
     */

    struct TrunkMeta {

        /*
         * Trunk Meta is used to holds information, which is stored in a shared
         * memory, to describe a Trunk object,
         * !!! This is not an actual object which a user shall interact with.
         */

        std::size_t block_size;
        std::size_t block_count;

        /* How many bytes have been used */
        std::atomic<std::size_t> used_block_count;

        /* How many bytes left in this Trunk */
        std::atomic<std::size_t> free_block_count;

        /* What is the max size of contiguous data byte */
        std::atomic<std::size_t> max_contiguous_blocks;

        /* Base on current object address, how many bytes is required to reach next trunk */
        std::atomic<std::size_t> next_trunk_shift;

        /* Base on current object address, how many bytes is required to step back to reach last trunk */
        std::atomic<std::size_t> last_trunk_shift;

        /* First usable block pointer shift */
        std::atomic<std::size_t> free_data;

        /* First used block pointer shift */
        ssize_t used_data;

        /* Constructor - only shall be called when the shared memory segment is initially applied.*/
        TrunkMeta *make_trunk(void *pBuffer, const std::size_t block_size, const std::size_t total_bytes);

        /* Constructor - used when other clients is attaching to it */
        TrunkMeta() = default;
    };

    TrunkMeta *TrunkMeta::make_trunk(void *__buff, const std::size_t __block_sz, const std::size_t __nbytes) {
        if (__buff == nullptr)
            throw nullptr;
        auto trunk_meta = static_cast<TrunkMeta *>(__buff);
        trunk_meta->block_size = __block_sz;
        /* according to total bytes and block size calculate block count*/
        trunk_meta->block_count = __nbytes / __block_sz;
        /* since this is freshly created trunk, max contiguous block should be block count*/
        trunk_meta->max_contiguous_blocks = trunk_meta->block_count;
        /* used block count should be 0 and free block block count should be trunk_meta->block_count */
        trunk_meta->used_block_count = 0;
        trunk_meta->free_block_count = trunk_meta->block_count;


    }

    class Trunk {
        /*
         * Trunk is a wrapper for Trunk Meta. When a segment of shared shared memory is
         * attach to a process, this will be constructed automatically, base on respectively
         * Trunk Meta.
         */

    private:
        /* Pointer to Trunk Meta in the shared memory */
        TrunkMeta *meta;

        /* Actual memory for allocation */
        [[nodiscard]] std::optional<void *> free_data() noexcept;

        /* Reserves one bit per block to indicate if it is free or used. */
        [[nodiscard]] std::optional<void *> used_data() noexcept;

        std::atomic<std::size_t> *max_contiguous_blocks;

        /* referring meta->used_bytes */
        std::atomic<std::size_t> *used_block_count;
        /* referring meta->free_bytes */
        std::atomic<std::size_t> *free_block_count;

        /* Pointer to next Trunk. !!! This is totally process dependent. */
        [[nodiscard]] Trunk *pNext() noexcept;

        /* Pointer to last Trunk. !!! This is totally process dependent. */
        [[nodiscard]] Trunk *pLast() noexcept;

        /* Try to find n of contiguous blocks and returns the first block's index or block_count on failure */
        std::optional<std::size_t> find_contiguous_block(std::size_t blocks) const noexcept;

        /* Marks n blocks as used starting at 'index' */
        void set_block_used(std::size_t index, std::size_t n = 1) noexcept;

        /* Mark n blocks as free starting at 'index' */
        void set_block_free(std::size_t index, std::size_t n = 1) noexcept;

        /* Pointer to head of this trunk */
        [[nodiscard]] std::byte *head() noexcept;


    public:

        std::size_t block_size() const;

        std::size_t block_count() const;

        /* Check if a pointer belongs to this trunk */
        bool belongs(void *ptr) const noexcept;

        /* Return nullptr if failed. */
        [[nodiscard]]
        void *allocate(std::size_t bytes) noexcept;

        /* free a segment of memory */
        void deallocate(void *ptr, std::size_t bytes) noexcept;

        Trunk(TrunkMeta *trunk_meta);
    };

    Trunk::Trunk(TrunkMeta *__meta) {
        if (__meta == nullptr)
            throw nullptr;
        this->free_block_count = &__meta->free_block_count;
        this->used_block_count = &__meta->used_block_count;
        this->max_contiguous_blocks = &__meta->max_contiguous_blocks;
    }

    std::optional<void *> Trunk::free_data() noexcept {
        return reinterpret_cast<std::byte *>(this->meta) + this->meta->free_data;
    }

    std::optional<void *> Trunk::used_data() noexcept {
        return reinterpret_cast<std::byte *>(this->meta) + this->meta->used_data;
    }

    std::byte *Trunk::head() noexcept {
        return reinterpret_cast<std::byte *>(const_cast<Trunk *>(this + sizeof(TrunkMeta)));
    }

    std::size_t Trunk::block_count() const {
        return static_cast<TrunkMeta *>(this->meta)->block_count;
    }

    std::size_t Trunk::block_size() const {
        return static_cast<TrunkMeta *>(this->meta)->block_size;
    }

    Trunk *Trunk::pNext() noexcept {
        return reinterpret_cast<Trunk *>(reinterpret_cast<std::byte *>(const_cast<Trunk *>(this)) +
                                         this->meta->next_trunk_shift);
    }

    Trunk *Trunk::pLast() noexcept {
        return reinterpret_cast<Trunk *>(reinterpret_cast<std::byte *>(const_cast<Trunk *>(this)) -
                                         this->meta->last_trunk_shift);
    }

    bool Trunk::belongs(void *__ptr) const noexcept {
        if (this <= __ptr &&
            __ptr <=
            reinterpret_cast<std::byte *>(const_cast<Trunk *>(this)) + this->block_size() * this->block_count())
            return true;
        return false;
    }

    void *Trunk::allocate(std::size_t __bytes) noexcept {
        /* Add 8 bytes, because the first 8 bytes will be used to stored size. */
        __bytes += 8;
        std::size_t required_blocks =
                __bytes % this->block_size() ? __bytes / this->block_size() + 1 : __bytes / this->block_size();
        if (required_blocks > *this->max_contiguous_blocks)
            return nullptr;
        auto found_index = this->find_contiguous_block(required_blocks);
        /* if the contiguous blocks are found */
        if (found_index.has_value()) {
            std::size_t *block_head = reinterpret_cast<std::size_t *>(this->head() +
                                                                      this->block_size() * found_index.value());
            // set block head to how many bytes
            *block_head = __bytes;
            // link it to used_data
            if (this->used_data().has_value())
                

        }
    }

}